import tensorflow as tf
import tensorflow_hub as hub
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
'''
print("++++++++++++",tf.executing_eagerly())
initial_model = keras.Sequential(
    [
        keras.Input(shape=(250, 250, 3)),
        layers.Conv2D(32, 5, strides=2, activation="relu"),
        layers.Conv2D(32, 3, activation="relu", name="my_intermediate_layer"),
        layers.Conv2D(32, 3, activation="relu"),
    ]
)
feature_extractor = keras.Model(
    inputs=initial_model.inputs,
    outputs=initial_model.get_layer(name="my_intermediate_layer").output,
)
# Call feature extractor on test input.
x = tf.ones((1, 250, 250, 3))
features = feature_extractor(x)
'''
print(tf.__version__)
import time
@tf.function
def qunat(value):
  #print("Tracing with", a)
  #value =  tf.function(a)
  a = time.time()
  max_range = tf.math.reduce_min(value)
  b = time.time()
  min_range = tf.math.reduce_min(value)
  c = time.time()
  result = tf.quantization.quantize(value,min_range,max_range,tf.quint8)
  d = time.time()
  return b-a,c-b,d-c


layer_name_list = ["input", "Conv2d_1a_3x3", "Conv2d_2a_3x3", "Conv2d_2b_3x3",
                   "MaxPool_3a_3x3", "Conv2d_3b_1x1", "Conv2d_4a_3x3",
                   "MaxPool_5a_3x3", "Mixed_5b", "Mixed_5c",
                   "Mixed_5d", "Mixed_6a", "Mixed_6b", "Mixed_6c",
                   "Mixed_6d", "Mixed_6e", "Mixed_7a",
                   "Mixed_7b", "Mixed_7c"]
model_name = "inception_v3"
result_pd = {}
for file_name in ["guitar.jpg","cat.jpg", "scorpion.jpg","dog.jpg", "pig.jpg"]:
    for layer_name in layer_name_list:
        data = np.asarray(np.load("middle_data/"+model_name+"/"+layer_name+"_"+file_name.split(".")[0]+".npy"),dtype=np.float32)
        result = []
        for i in range(30):
            temp = qunat(data)
            result.append(temp)
        print(file_name,layer_name,np.around(np.average(np.array(result)),3))
